# Tests the infinite dirichlet distribution

from pyblog import *
from testmain import MyTestCase, testmodule
from random import Random

class test_infdir_gmm(MyTestCase):
  def test_mean(self):
    """
    The first element drawn from an infinite dirichlet-categorical must
    have expected probability 1/(1+alpha)
    """
    alpha = 9
    inf = InfDirichlet(alpha)
    rand = Random()
    first_prob = []
    for j in range(1000):
      wts = inf.sample(rand)
      cat = InfCategorical(wts)
      first = cat.sample(rand)
      first_cnt = 1 + sum(first == cat.sample(rand) for i in range(100))
      first_prob.append(float(first_cnt) / 100)

    avg_first_prob = sum(first_prob)/len(first_prob)

    if dispstats:
      print "avg. first wt.", avg_first_prob

    self.assertAbsEqual(avg_first_prob, 0.1, .1, "avg. first probability")

  def test_conjugacy(self):
    from pyblog.misc.poly import Polynomial
    x = Polynomial([.1,.2,.3])

    stats = InfCategorical.SufficientStats()
    for v in [0,0,0,1,2,1,0,3]:
      stats.add(v)
    like = InfCategorical(x)(stats)

    post = like * InfDirichlet(5)

    if dispstats:
      print "likelihood:", str(like)
      print "posterior:", str(post)

    rand = Random()
    prob_0 = sum(InfCategorical(post.sample(rand)).sample(rand) == 0 \
                 for i in range(10000)) / 10000.0

    if dispstats:
      print "0 prob", prob_0

    self.assertAbsEqual(prob_0, 4.0/13, .1, "probability of drawing 0")
    
  def testgmm(self):
    
    @var_dist
    def cluster_wts(): return InfDirichlet(5)
    
    @var_dist
    def cluster(i): return InfCategorical(cluster_wts())
    
    @var_dist
    def mean(i): return Normal(0, .001)
    
    @var_dist
    def prec(i): return Gamma(1, 1)
    
    @var_dist
    def coord(i): return Normal(mean(cluster(i)), prec(cluster(i)))

    @var
    def count_clusters(n):
      return len(set(cluster(i) for i in range(n)))
    
    # sample some values
    N = 10
    data = query([coord(i) for i in range(N)],[], scans=0, burnin=0,
                 outtyp=QUERY_LAST, stats=False)

    # infer the means
    @var
    def all_means(n):
      means = [(mean(k), prec(k), cluster_wts()[k])\
               for k in set(cluster(i) for i in range(n))]
      means.sort()
      return means
    
    means = query([all_means(len(data))],
                  [coord(i)==c for i,c in enumerate(data)],
                  burnin=0, scans=10,
                  stats=dispstats, outtyp=QUERY_LAST)
    data.sort()
    if dispstats:
      print "data", data
      print "means", means
    
if __name__ == "__main__":
  testmodule(__import__("testinfdir"))
